import json
from typing import AsyncGenerator, Dict
import logging

logger = logging.getLogger(__name__)

# ------------------------------
# 输入文本替换：用户提问预处理
# ------------------------------
def apply_text_replacements(text: str) -> str:
    """
    应用文本替换规则的通用函数（用户输入预处理）
    可在此处添加更多替换规则，便于统一管理
    """
    if not text:
        return text

    # 定义输入替换规则字典：key=待替换文本，value=目标文本
    input_replacement_rules: Dict[str, str] = {
        "你": "你",
        # 后续新增输入替换规则直接加在这里
        # "你们": "日落前的星星们",
        # "你的": "日落前的星星的",
    }

    # 批量应用所有替换规则
    for old_str, new_str in input_replacement_rules.items():
        text = text.replace(old_str, new_str)

    return text


# ------------------------------
# 输出文本替换：智能体响应后处理
# ------------------------------
def get_response_replacement_rules() -> Dict[str, str]:
    """
    定义输出替换规则字典（集中管理，便于维护）
    后续新增输出替换规则直接在此函数中添加
    """
    return {
        "日落前的星星": "我",
        "kaelis": "我",
        # 新增规则示例：
        # "kaelis": "我",
        # "它": "我",
    }


async def replace_response_content(original_stream: AsyncGenerator[str, None]) -> AsyncGenerator[str, None]:
    """
    拦截原始流式响应，应用输出替换规则后返回给前端
    替换规则通过 get_response_replacement_rules() 统一获取
    """
    # 提前加载输出替换规则（避免每次迭代都重新创建）
    response_rules = get_response_replacement_rules()

    async for chunk in original_stream:
        # 仅处理 "data:" 开头的流式片段
        if chunk.startswith("data:"):
            json_str = chunk.lstrip("data:").strip("\n")
            try:
                data_dict = json.loads(json_str)

                # 只处理 streaming 类型的业务数据（跳过 [DONE] 等标记）
                if data_dict.get("msg") == "streaming" and "data" in data_dict:
                    modified_data = data_dict["data"]
                    # 批量应用所有输出替换规则
                    for old_str, new_str in response_rules.items():
                        modified_data = modified_data.replace(old_str, new_str)
                    # 更新数据字典并重新生成 JSON
                    data_dict["data"] = modified_data
                    modified_json = json.dumps(data_dict, ensure_ascii=False)
                    yield f"data:{modified_json}\n\n"
                else:
                    # 非业务数据（如结束标记）直接返回原内容
                    yield chunk

            except json.JSONDecodeError:
                logging.warning(f"流式响应JSON解析失败，跳过替换: {chunk}")
                yield chunk
        else:
            # 非 "data:" 格式的片段直接返回
            yield chunk